#!/usr/bin/python

# -*- coding: utf-8 -*-

from gevent import monkey

monkey.patch_all()
from gevent import socket
from gevent.pool import Pool
import traceback
import BaseHTTPServer
#from gevent import socket
from gevent.server import StreamServer
from Crypto.Cipher import ARC4
from gevent.select import select
import urlparse
#import zlib
import httplib
import StringIO
import mimetools
import util
import ConfigParser

proxy_host = "127.0.0.1"
proxy_port =  8038

headers_to_client={
    "Connection":"keep-alive",
    "Accept-Ranges": "bytes",
    "Content-Type": "application/json",
    #"Transfer-Encoding":"deflate",
    #"Content-Encoding":"deflate",
}


class Handler(BaseHTTPServer.BaseHTTPRequestHandler):
    server_version = "pws"
    protocol_version = "HTTP/1.1"
    def __init__(self, request, addr):
        self.request = request
        self.client_address = addr
        #self.rfile = request.makefile("rb", -1)
        #self.wfile = request.makefile("wb", 0)
        self.setup()
        self.handle()
        self.finish()

    def get_host(self, url):
        if url.startswith("http"):
            port = 80
            parsed = urlparse.urlparse(url)
            host = parsed.hostname
            if parsed.port:
                port = parsed.port
            return host,int(port)
        elif ":" in url:    
            host,port = url.split(":", 1)
            return host,int(port)
        raise TypeError("invalid url")
        
    def get_request_path(self, url):
        parsed = urlparse.urlparse(url)
        path = "/"
        if parsed.path:
            path = parsed.path
        if parsed.query:
            path += "?%s" % parsed.query
        return path
        
    def do_GET(self):
        if self.path.startswith("/"):
            msg = "hello, world"
            self.send_response(200)
            self.send_header("Content-Length", len(msg))
            self.log_request(200, len(msg))
            self.send_header("Connection", "close")
            self.end_headers()
            self.wfile.write(msg)
            return
        
        self.send_response(200)
        self.send_header("Connection","close")
        self.end_headers()
        self.log_request(200)
        
    def do_POST(self):
        #print self.headers
        body_len = int(self.headers.get("content-length", 0))
        body = ''
        if body_len:
            body = self.rfile.read(body_len)
            
        transfer_encoding = self.headers.get("transfer-encoding", '')
        #print transfer_encoding    
        #print "POST %s with body '%s'" % (self.path, body)
        if not body and not transfer_encoding:
            self.send_response(400)
            self.send_header("Connection", "close")
            self.end_headers()
            return
        sid = self.headers.get("Cookie").split("=",1)[1]
        #print sid
        self.rkey = rkey = util.decode_key(sid, password)
        self.renc = enc = ARC4.new(rkey)
        
        if transfer_encoding:
            l = self.rfile.readline().strip(" \r\n")
            l = int(l, 16)
            unfinished = l
            body = ''
            while unfinished > 0:
                d = self.rfile.read(unfinished)
                if not d:
                    break
                unfinished -= len(d)
                body += d
            self.rfile.readline()
        data = enc.encrypt(body)
        #print len(data)
        fp = StringIO.StringIO(data)
        firstline = fp.readline(65535)
        #print firstline

        words = firstline.strip(" \r\n").split()

        # may be password is wrong
        if len(words) != 3:
            self.send_response(400)
            self.send_header("Connection","close")
            self.end_headers()
            return

        if words[0] not in ["GET","POST","CONNECT"]:
            self.send_response(400)
            self.send_header("Connection","close")
            self.end_headers()
            return
        if words[2].lower() not in ["http/1.1","http/1.0"]:
            self.send_response(400)
            self.send_header("Connection","close")
            self.end_headers()
            return

        self.request_method,self.request_url,self.request_version = words

        self.request_headers = mimetools.Message(fp)
        l = int(self.request_headers.get("Content-Length",0))
        self.request_body = ''
        if l:
            self.request_body = fp.read(l)
            fp.close()
        if self.request_method != "CONNECT":
            self.do_httprelay()
        else:
            self.do_forward()

    def do_forward(self):
        host,port = self.get_host(self.request_url)
        #print "connecting to %s:%d" % (host, port)
        self.rsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        try:
            self.rsock.connect((host, port))
            #print "connect to %s:%d success" % (host,port)
        except:
            traceback.print_exc()
            self.log_message("%s %s %s -", self.request_method, self.request_url, "503")
            self.send_response(503)
            self.send_header("Connection","close")
            self.end_headers()
            return


        key, sessionid = util.random_key(password)
        enc = ARC4.new(key)
        #self.wfile.write("HTTP/1.1 200 OK\r\n")
        self.send_response(200)
        for k,v in headers_to_client.items():
            #self.wfile.write("%s: %s\r\n" % (k.title(),v))
            self.send_header(k, v)
        #self.wfile.write("Set-Cookie: sessionid=%s\r\n" % sessionid))
        self.send_header("Transfer-Encoding", "chunked")
        self.send_header("Content-Encoding", "deflate")
        self.send_header("Set-Cookie", "sessionid=%s" % sessionid)
        #self.wfile.write("\r\n")
        self.end_headers()
        
        self.log_message("%s %s %s -", self.request_method, self.request_url, "200")
        
        reply = "HTTP/1.1 200 OK\r\n\r\n"
        enc_reply = enc.encrypt(reply)
        self.wfile.write("%x\r\n%s\r\n" % (len(enc_reply), enc_reply))
        #print "begin to forward"
        while True:
            r,_,_ = select([self.request, self.rsock],[],[])
            if self.request in r:
                l = self.rfile.readline().strip(" \r\n")
                if not l:
                    break
                l = int(l, 16)
                data = ''
                unfinished = l
                while unfinished > 0:
                    d = self.rfile.read(unfinished)
                    unfinished -= len(d)
                    data += d
                self.rfile.readline()

                self.rsock.send(self.renc.encrypt(data))
            if self.rsock in r:
                data = self.rsock.recv(1024)
                if not data:
                    break
                enc_data = enc.encrypt(data)
                self.wfile.write("%x\r\n%s\r\n" % (len(enc_data),enc_data))
        self.rsock.close()

    def do_httprelay(self):
        req_path = self.get_request_path(self.request_url)
        host,port = self.get_host(self.request_url)
        conn = httplib.HTTPConnection(host, port)
        
        # copy headers
        h2 ={}
        for k,v in self.request_headers.items():
            h2[k.title()] = v

        try:
            conn.request(self.request_method, req_path, self.request_body, h2)
            response = conn.getresponse()
            body = response.read()
        except:
            traceback.print_exc()
            self.send_response(500)
            self.send_header("Connection","close")
            self.end_headers()
            #self.log_request(500)
            self.log_message('%s %s" %d %s' % (self.request_method, self.request_url, 500, "-"))
            conn.close()
            return

        conn.close()
        data_to_send  = "%s %s %s\r\n" % (self.request_version, response.status, response.reason)
        #print data_to_send
        #print response.getheaders()
        for k,v in response.getheaders():
            if k.lower() == 'transfer-encoding' or k.lower() == 'content-length':
                continue
            data_to_send += '%s: %s\r\n' % (k.title(),v)
            
        data_to_send += 'Content-Length: %d\r\n' % len(body) 
        data_to_send += "Proxy-Connection: keep-alive\r\n"
        #print data_to_send
        data_to_send += '\r\n'
        data_to_send += body
        lkey,sessionid = util.random_key(password)
        enc = ARC4.new(lkey)
        enc_data = enc.encrypt(data_to_send)
        self.send_response(200)
        
        for k,v in headers_to_client.items():
            self.send_header(k,v)
            
        self.send_header("Content-Length", len(enc_data))   
        self.send_header("Set-Cookie", "sessionid=%s" % sessionid)
        self.end_headers()
        self.wfile.write(enc_data)
        #self.log_request(200, len(enc_data))
        self.log_message("%s %s %d %d", str(self.request_method), str(self.request_url), 200, len(enc_data))

    def address_string(self):
        return "%s:%d" % self.client_address

        
    def do_CONNECT(self):
        self.send_response(400)
        self.send_header("Connection", "close")
        self.end_headers()

    def send_response(self, code, message=None):
        if message is None:
            if code in self.responses:
                message = self.responses[code][0]
            else:
                message = ''
        self.wfile.write("%s %d %s\r\n" % (
            self.protocol_version, code, message))
        self.send_header("Server", self.server_version)
        self.send_header("Date", self.date_time_string())

    def send_header(self, keyword, value, fp=None):
        out = self.wfile
        if fp:
            out = fp
        out.write("%s: %s\r\n" % (keyword.title(), value))
        if keyword.lower() == 'connection' or keyword.lower() == 'proxy-connection':
            if value.lower() == 'close':
                self.close_connection = 1
            elif value.lower() == 'keep-alive':
                self.close_connection = 0

    def end_headers(self, fp=None):
        out = self.wfile
        if fp:
            out = fp
        out.write("\r\n")


if __name__ == "__main__":
    cfg = ConfigParser.ConfigParser()
    cfg.read("server.ini")
    host = cfg.get("server","listen")
    port = cfg.getint("server","port")
    password = cfg.get("server","password")
    pool = Pool(100)
    srv = StreamServer((host, port), Handler, spawn = pool)
    srv.serve_forever()

